from __future__ import annotations

"""
调度器服务
----------
E1 阶段：将 D1→D4 按固定节奏顺序执行，且串行避免重叠
"""

import asyncio
import logging
import random
from contextlib import suppress
from datetime import datetime, timezone
from typing import Optional

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from ..config import get_settings
from ..db import SessionLocal
from ..models import Decision, Run, RunMode, Snapshot
from .trading_pipeline import (
    create_decision,
    create_snapshot,
    execute_decision,
    validate_decision,
)

logger = logging.getLogger(__name__)
settings = get_settings()


class TradingScheduler:
    """交易执行调度器"""
    
    def __init__(
        self,
        run_id: str,
        exchange: str,
        symbol: str,
        account_name: str,
        risk_config: dict,
        interval_seconds: int = 60,
        timeout_per_stage: dict = None,
        max_retries_per_stage: int = 3,
        jitter_seconds: int = 2,
    ):
        self.run_id = run_id
        self.exchange = exchange
        self.symbol = symbol
        self.account_name = account_name
        self.risk_config = risk_config
        self.interval_seconds = interval_seconds
        self.timeout_per_stage = timeout_per_stage or {
            "snapshot": 5,
            "decision": 60,
            "risk_gate": 10,
            "execute": 10,
        }
        self.max_retries_per_stage = max_retries_per_stage
        self.jitter_seconds = jitter_seconds
        self.running = False
        self.task: Optional[asyncio.Task] = None
    
    async def _run_stage_with_retry(self, stage_name: str, stage_func, *args):
        """执行阶段，带重试机制（指数退避）"""
        session_arg = next(
            (arg for arg in args if isinstance(arg, AsyncSession)),
            None,
        )

        for attempt in range(self.max_retries_per_stage):
            task = asyncio.create_task(stage_func(*args))
            try:
                return await asyncio.wait_for(
                    task,
                    timeout=self.timeout_per_stage.get(stage_name, 10),
                )
            except asyncio.TimeoutError:
                logger.warning(
                    f"{stage_name} 超时 (尝试 {attempt + 1}/{self.max_retries_per_stage})"
                )
                task.cancel()
                with suppress(asyncio.CancelledError, Exception):
                    await task

                if session_arg is not None:
                    with suppress(Exception):
                        await session_arg.rollback()

                if attempt < self.max_retries_per_stage - 1:
                    wait_time = 2 ** attempt  # 指数退避
                    await asyncio.sleep(wait_time)
                else:
                    logger.error(f"{stage_name} 失败，已达到最大重试次数")
                    raise
            except Exception as e:
                logger.error(
                    f"{stage_name} 失败: {e} (尝试 {attempt + 1}/{self.max_retries_per_stage})"
                )

                if session_arg is not None:
                    with suppress(Exception):
                        await session_arg.rollback()

                if attempt < self.max_retries_per_stage - 1:
                    wait_time = 2 ** attempt
                    await asyncio.sleep(wait_time)
                else:
                    raise
        return None
    
    async def _execute_tick(self) -> None:
        """执行一个完整的 tick：快照 → 决策 → 风控 → 执行"""
        async with SessionLocal() as session:
            try:
                # 获取 run 信息
                stmt = select(Run).where(Run.run_id == self.run_id)
                result = await session.execute(stmt)
                run = result.scalar_one_or_none()
                if not run:
                    logger.error(f"Run 不存在: {self.run_id}")
                    return
                
                mode = run.mode
                # 确保 mode 是 RunMode 枚举或字符串
                if isinstance(mode, RunMode):
                    mode_enum = mode
                    mode_str = mode.value
                else:
                    mode_enum = RunMode(mode)  # 如果是字符串，转换为枚举
                    mode_str = mode
                
                # 验证调度器配置（如果为空，使用 run 中的配置）
                exchange = self.exchange or run.exchange
                symbol = self.symbol or run.symbol
                account_name = self.account_name or run.account_name
                
                if not exchange or not symbol or not account_name:
                    logger.error(
                        f"调度器配置不完整: run_id={self.run_id}, "
                        f"exchange={exchange}, symbol={symbol}, account_name={account_name}"
                    )
                    return
                
                # D1: 快照
                logger.info(f"[Tick] 开始快照: run_id={self.run_id}, exchange={exchange}, symbol={symbol}")
                snapshot = await self._run_stage_with_retry(
                    "snapshot",
                    create_snapshot,
                    session,
                    self.run_id,
                    exchange,
                    symbol,
                    account_name,
                    self.risk_config,
                )
                
                # D2: 决策
                logger.info(f"[Tick] 开始决策: run_id={self.run_id}")
                # 在传递给 create_decision 之前，先提取 snapshot 的所有属性值
                # 这样即使在超时重试时，这些值也是可用的（不依赖于已 rollback 的 session）
                snapshot_exchange = snapshot.exchange
                snapshot_symbol = snapshot.symbol
                snapshot_payload = snapshot.payload if snapshot.payload else {}
                snapshot_id = snapshot.id
                
                context = {
                    "account_name": account_name,
                    "run_id": self.run_id,
                    "timezone": "Asia/Shanghai",
                    "mode": mode_str,
                }
                decision = await self._run_stage_with_retry(
                    "decision",
                    create_decision,
                    session,
                    self.run_id,
                    snapshot_id,
                    snapshot_exchange,
                    snapshot_symbol,
                    snapshot_payload,
                    self.risk_config,
                    context,
                )
                
                # D3: 风控闸门
                logger.info(f"[Tick] 开始风控校验: decision_id={decision.id}")
                decision = await self._run_stage_with_retry(
                    "risk_gate",
                    validate_decision,
                    session,
                    decision,
                    snapshot,
                    self.risk_config,
                )
                
                # D4: 执行器（如果通过风控）
                if decision.status.value == "APPROVED":
                    logger.info(f"[Tick] 开始执行: decision_id={decision.id}")
                    await self._run_stage_with_retry(
                        "execute",
                        execute_decision,
                        session,
                        decision,
                        mode_enum,  # 传递枚举类型
                        snapshot,  # 传递快照用于生成 client_order_id
                        30,  # 冷却窗：30秒
                    )
                else:
                    logger.info(f"[Tick] 决策未通过风控: reason={decision.reason}")
                
                logger.info(f"[Tick] 完成: run_id={self.run_id}")
                
            except Exception as e:
                logger.error(f"[Tick] 执行失败: {e}", exc_info=True)
    
    async def start(self) -> None:
        """启动调度器"""
        if self.running:
            logger.warning("调度器已在运行")
            return
        
        self.running = True
        self.task = asyncio.create_task(self._run_loop())
        logger.info(f"调度器已启动: run_id={self.run_id}, interval={self.interval_seconds}s")
    
    async def stop(self) -> None:
        """停止调度器"""
        self.running = False
        if self.task:
            self.task.cancel()
            try:
                await self.task
            except asyncio.CancelledError:
                pass
        logger.info(f"调度器已停止: run_id={self.run_id}")
    
    async def _run_loop(self) -> None:
        """主循环"""
        while self.running:
            try:
                await self._execute_tick()
                
                # 等待下一个 tick（加入 jitter）
                jitter = random.uniform(0, self.jitter_seconds)
                await asyncio.sleep(self.interval_seconds + jitter)
                
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"调度循环错误: {e}", exc_info=True)
                await asyncio.sleep(5)  # 错误后短暂等待


# 全局调度器实例（按 run_id 管理）
_schedulers: dict[str, TradingScheduler] = {}


async def start_scheduler(
    run_id: str,
    exchange: str,
    symbol: str,
    account_name: str,
    risk_config: dict,
    mode: RunMode = RunMode.shadow,
    interval_seconds: int = 60,
) -> TradingScheduler:
    """启动一个新的调度器"""
    # 创建或获取 run，并保存调度器配置（用于重启后恢复）
    async with SessionLocal() as session:
        from ..config import get_settings
        settings = get_settings()
        
        stmt = select(Run).where(Run.run_id == run_id)
        result = await session.execute(stmt)
        run = result.scalar_one_or_none()
        
        if not run:
            run = Run(
                run_id=run_id,
                model=settings.deepseek_model or "deepseek-reasoner",
                mode=mode,
                exchange=exchange,
                symbol=symbol,
                account_name=account_name,
                risk_config=risk_config,
                interval_seconds=interval_seconds,
            )
            session.add(run)
            await session.commit()
        else:
            # 更新调度器配置（允许修改）
            run.exchange = exchange
            run.symbol = symbol
            run.account_name = account_name
            run.risk_config = risk_config
            run.interval_seconds = interval_seconds
            run.mode = mode
            await session.commit()
        
        scheduler = TradingScheduler(
            run_id=run_id,
            exchange=exchange,
            symbol=symbol,
            account_name=account_name,
            risk_config=risk_config,
            interval_seconds=interval_seconds,
        )
        
        _schedulers[run_id] = scheduler
        await scheduler.start()
        
        return scheduler


async def stop_scheduler(run_id: str) -> None:
    """停止调度器"""
    if run_id in _schedulers:
        await _schedulers[run_id].stop()
        del _schedulers[run_id]


async def stop_all_schedulers() -> None:
    """停止所有调度器"""
    for scheduler in list(_schedulers.values()):
        await scheduler.stop()
    _schedulers.clear()

